package com.gzist.elderlyserviceplatform.controller.common;

import cn.hutool.core.util.StrUtil;
import com.alibaba.cloud.ai.advisor.RetrievalRerankAdvisor;
import com.alibaba.cloud.ai.model.RerankModel;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.client.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
import org.springframework.ai.rag.generation.augmentation.QueryAugmenter;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.web.bind.annotation.*;

import java.util.Map;

/**
 * 向量数据库接口层
 */
@Tag(name = "向量数据库接口层")
@RestController
@RequestMapping("/api/common/rag")
@RequiredArgsConstructor
public class MilvusController {

    private final VectorStore vectorStore;
    private final ChatModel chatModel;
    private final RerankModel rerankModel;

    // 系统提示词模板（角色定义）
    private static final String SYSTEM_PROMPT = """
            你是一个专业的养老领域知识助手，具备以下能力：
            1. 严格基于提供的上下文信息回答问题
            2. 对复杂查询进行多维度分析
            3. 保持回答的结构化和可追溯性
            4. 发现数据异常时主动提示
            """;

    // 用户提示词模板（重排序阶段）
    private static String USER_PROMPT_RERANK = """
            以下是知识库检索信息：
            --------------------
            {question_answer_context}
            --------------------

            <question>
            用户问题：{question}
            </question>
            请基于上述信息，结合历史对话内容回答用户问题。回答应：
            1. 内容完整且紧扣问题核心
            2. 避免直接摘抄原文，需语义转化
            3. 语言简洁专业，逻辑清晰连贯
            4. 若信息不足请礼貌告知用户

            回答建议：先直接回应问题关键，再展开必要解释，确保用户获得准确且易于理解的信息。
            """;

    @GetMapping("/retrieval")
    @Operation(summary = "基础检索",description = "一般检索，没有检索前处理和检索后处理")
    public String retrieval(String question) {
        // 配置检索顾问 其内置了用户模板，无需手动传入
        RetrievalAugmentationAdvisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder()
                .documentRetriever(VectorStoreDocumentRetriever.builder()
                        .similarityThreshold(0.2)   // 设置相似度阈值
                        .topK(7)                    // 设置返回前7条相关结果
                        .vectorStore(vectorStore)   // 设置向量存储
                        .build())
                .build();

        ChatClient chatClient = ChatClient.builder(chatModel)
                .defaultSystem(SYSTEM_PROMPT)  // 设置系统角色
                .defaultAdvisors(
                        // 基础检索顾问
                        retrievalAugmentationAdvisor
                )
                .build();

        return chatClient.prompt()
                .user(question)
                .call()
                .content();
    }

    @GetMapping("/retrievalByRerank")
    @Operation(summary = "检索后优化",description = "基础检索+重排序")
    public String retrievalByRerank(String question) {
        // 因为使用PromptTemplate转换模板需要对所有的变量名进行注入值，而在这里对于检索结果是无需注入的，
        // 因此无法使用PromptTemplate进行转换，如果要添加其它变量在模板中，在这是使用字符串拼接的方式
        USER_PROMPT_RERANK = StrUtil.format(USER_PROMPT_RERANK, Map.of("question", question));
        // 配置搜索参数
        SearchRequest searchRequest = SearchRequest.builder()
                .similarityThreshold(0.2)
                .topK(7)
                .build();
        // 配置检索顾问
        RetrievalAugmentationAdvisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder()
                .documentRetriever(VectorStoreDocumentRetriever.builder()
                        .similarityThreshold(0.2)   // 设置相似度阈值
                        .topK(7)                    // 设置返回前7条相关结果
                        .vectorStore(vectorStore)   // 设置向量存储
                        .build())
                .order(1)                           // 指定执行顺序(先执行检索前处理与检索)
                .build();
        // 配置重排序顾问
        RetrievalRerankAdvisor retrievalRerankAdvisor = new RetrievalRerankAdvisor(
                vectorStore,
                rerankModel,
                searchRequest,
                // 创建带变量注入的模板文本(只能存在一个变量即最终的检索结果)
                USER_PROMPT_RERANK,
                0.6  // 重排序阈值
        );
        // 构建ChatClient（包含双阶段RAG处理）
        ChatClient chatClient = ChatClient.builder(chatModel)
                .defaultSystem(SYSTEM_PROMPT)  // 设置系统角色
                .defaultAdvisors(
                        // 基础检索顾问（使用第一阶段模板）
                        retrievalAugmentationAdvisor,
                        // 重排序顾问（使用优化后模板）
                        retrievalRerankAdvisor
                )
                .build();

        return chatClient.prompt()
                .user(question)
                .call()
                .content();
    }
}